
import sys

KNOWN_GENE_TABLE = "/home/bowtie/Bowtie_Illumina_Alignment/references/mergedGeneCluster-plus-additional.table.txt"
CHR_LENGTHS = "chromosome_lengths"

def intervalsOverlap(ival1, ival2):
    (s1, e1) = ival1
    (s2, e2) = ival2
    return (e1>=s2 and s1<=e2)

def findGenes(query, filename=KNOWN_GENE_TABLE):
    (chr, strand, spos, epos) = query
    genes = {}
    queryIval = (spos, epos)
    for line in open(filename):
        if line[0]=="#": continue
        cols = line.strip().split("\t")
        (geneId, geneChr, geneStrand, geneStart, geneEnd) = tuple(cols[:5])
        if chr!=geneChr: continue
        if strand!=geneStrand and strand!='?': continue
        geneIval = (long(geneStart), long(geneEnd))
        if intervalsOverlap(queryIval, geneIval):
            genes[geneId] = geneIval
    return genes

def findExonIntrons(query, filename=KNOWN_GENE_TABLE):
    (chr, strand, spos, epos) = query
    exonIntrons = []
    queryIval = (spos, epos)
    for line in open(filename):
        if line[0]=="#": continue
        cols = line.strip().split("\t")
        (geneId, geneChr, geneStrand, geneStart, geneEnd, geneFieldA, geneFieldB, geneExonCount, geneExonStarts, geneExonEnds) = tuple(cols[:10])
        if chr!=geneChr: continue
        if strand!=geneStrand and strand!='?': continue
        geneIval = (long(geneStart), long(geneEnd))
        if not intervalsOverlap(queryIval, geneIval): continue
        exonStarts = [long(p) for p in geneExonStarts.split(",")[:-1]]
        exonEnds = [long(p) for p in geneExonEnds.split(",")[:-1]]
        for i in xrange(0, int(geneExonCount)):
            exonIntrons.append(((exonStarts[i], exonEnds[i]), "Exon"))
            if i>0 and exonStarts[i]-1>=exonEnds[i-1]+1:
                exonIntrons.append(((exonEnds[i-1]+1, exonStarts[i]-1), "Intron"))
    return exonIntrons

class Label:
    def __init__(self, labels={}):
        self.labels = labels.copy()
    def add(self, label):
        self.labels[label] = True
        return self
    def copy(self):
        return Label(self.labels)
    def __str__(self):
        if not self.labels:
            return "Unlabeled"
        return str(self.labels.keys())

class LabeledIntervalList:
    def __init__(self, ll):
        self.ll = ll
        self.ivals = [((1, ll), Label())]

    def getLeftIntersection(self, ival):
        """return the smallest index ix s.t. ivals[ix]'s end >= ival's start"""
        mn=0
        (mnIval, mnLabel) = self.ivals[0]
        if mnIval[1] >= ival[0]: return 0
        mx=len(self.ivals)
        while mx-mn>1:
            md = (mn+mx)/2
            (mdIval, mdLabel) = self.ivals[md]
            if mdIval[1] >= ival[0]:
                mx=md
            else:
                mn=md
        return mx

    def getLabels(self, ival):
        ix = self.getLeftIntersection(ival)
        labels = {}
        while ix<len(self.ivals):
            (oldIval, oldLabel) = self.ivals[ix]
            if ival[1]<oldIval[0]:
                break
            if oldLabel == None:
                continue
            for label in oldLabel.labels.keys():
                labels[label] = True
            ix += 1
        return labels               

    def addLabel(self, ival, label):
        ix = self.getLeftIntersection(ival)
        newIvals = []
        if ix>0: newIvals = self.ivals[:ix]
        while ix<len(self.ivals):
            (oldIval, oldLabel) = self.ivals[ix]
            #if ival[0]>oldIval[1]:
            #    newIvals.append((oldIval, oldLabel))
            #    ix+=1
            #    continue
            if ival[1]<oldIval[0]:
                newIvals.extend(self.ivals[ix:])
                break
            if oldIval[0]<ival[0]:
                # start of old interval is unchanged
                newIvals.append(((oldIval[0], ival[0]-1), oldLabel.copy()))
            if oldLabel == None:
                # adding to previously unlabeled region
                newLabel = Label(label)
            else:
                # adding to existing labels
                newLabel = oldLabel.copy().add(label)
            # new label attachs to intersection
            newIvals.append(((max(oldIval[0],ival[0]), min(oldIval[1], ival[1])), newLabel))
            if oldIval[1]>ival[1]:
                # end of old interval is unchanged
                newIvals.append(((ival[1]+1, oldIval[1]), oldLabel))
            ix+=1
        self.ivals = newIvals

def getAllGeneLabels(extender = 0):
    allLabels = {}
    for line in open(CHR_LENGTHS):
        (ll, chr) = line.strip().split(" ")
        ll = long(ll)
        chr = chr.replace(".fa", "")
        genes = findGenes((chr, "?", 1, ll))
        ivals = LabeledIntervalList(ll)
        for gene in genes.keys():
            (spos, epos) = genes[gene]
            ivals.addLabel((spos - extender, epos + extender), gene)
        allLabels[chr] = ivals
        print >> sys.stderr, "# Read ", len(ivals.ivals), " intervals from ", chr
    return allLabels

if __name__=="__main__":
    for line in open(CHR_LENGTHS):
        (ll, chr) = line.strip().split(" ")
        ll = long(ll)
        chr = chr.replace(".fa", "")
        if len(sys.argv)>1 and sys.argv[1]=="-E":
            ei = findExonIntrons((chr, "?", 1, ll))
            print >> sys.stderr, "got %ld exon/intron labels on %s" % (len(ei), chr)
            ivals = LabeledIntervalList(ll)
            for (ival, label) in ei:
                ivals.addLabel(ival, label)
                if len(ivals.ivals)%10000<10: print >> sys.stderr, "ivals length = %ld" % len(ivals.ivals)
        else:
            genes = findGenes((chr, "?", 1, ll))
            ivals = LabeledIntervalList(ll)
            for gene in genes.keys():
                (spos, epos) = genes[gene]
                ivals.addLabel((spos - 5000, epos + 5000), gene)
        prevLabels = None
        prevIval = None
        for (ival, label) in ivals.ivals:
            if not label.labels:
                if prevLabels and ival[1]-prevIval[1] < 5000:
                    prevIval[1] = ival[1]
                    continue
                if prevLabels:
                    print "%s\t%d\t%d\t%d\t%s" % (chr, prevIval[0], prevIval[1], prevIval[1]-prevIval[0]+1, str(prevLabels.keys()))
                print "%s\t%d\t%d\t%d\t%s" % (chr, ival[0], ival[1], ival[1]-ival[0]+1, label)
                prevIval = [ival[1]+1, ival[1]+1]
                prevLabels = {}
            else:
                if not prevLabels:
                    prevLabels = {}
                    prevIval = [ival[0], ival[0]]
                prevIval[1] = ival[1]
                for l in label.labels:
                    prevLabels[l] = True
        if prevLabels:
            print "%s\t%d\t%d\t%d\t%s" % (chr, prevIval[0], prevIval[1], prevIval[1]-prevIval[0]+1, str(prevLabels.keys()))
